#pragma once
#ifdef _WIN32
#pragma comment(lib, "ws2_32.lib")
#endif
#include <iostream>
#include <cstring>
#include <cerrno>
#include <iostream>
#include <cstring>
#include <cerrno>
#include <stdlib.h>
#ifdef _WIN32
#include <winsock2.h>
#include <windows.h>
#elif __linux__
#include <sys/types.h>
#include <sys/fcntl.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/fcntl.h>
#include <unistd.h>
#endif
#define TCP_SOCKET SOCK_STREAM
#define UDP_SOCKET SOCK_DGRAM
#ifdef _WIN32
typedef int socketlen_t;
#elif __linux__
typedef socklen_t socketlen_t;
#endif

class socket_t
{
public:
    socket_t(const uint16_t port, const uint8_t type)
        : _port(port), _fd(-1)
    {
#ifdef _WIN32
        WSADATA wsa_data;
        (void)WSAStartup(MAKEWORD(2, 2), &wsa_data);
#endif
        // socket
        _fd = socket(AF_INET, type, 0);
        if (_fd < 0)
        {
            std::cerr << strerror(errno) << std::endl;
            exit(1);
        }
        // info
        sockaddr_in my_addr;
        memset(&my_addr, 0, sizeof(my_addr));
        my_addr.sin_family = AF_INET;
        my_addr.sin_addr.s_addr = INADDR_ANY;
        my_addr.sin_port = htons(_port);
        // may use the port again immediately after restrating the socket
        set_sock_reuse();
        // bind
        if (bind(_fd, (sockaddr *)&my_addr, sizeof(sockaddr_in)) == -1)
        {
            std::cerr << strerror(errno) << std::endl;
            exit(2);
        }
    }
#ifdef __linux__
    bool set_sock_nonblock()
    {
        // 设置listen_fd 非阻塞
        int falg = fcntl(_fd, F_GETFL);
        return fcntl(_fd, F_SETFL, falg | O_NONBLOCK) != -1;
    }
#endif
    virtual ~socket_t()
    {
#ifdef _WIN32
        closesocket(_fd);
#elif __linux__
        close(_fd);
#endif
    }
    // get fd
    int get_fd() const
    {
        return _fd;
    }
    // get port
    int get_port() const
    {
        return _port;
    }

private:
    void set_sock_reuse()
    {
        // opt should be int
        int opt = 1;
        setsockopt(_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof opt);
#ifdef __linux__
        setsockopt(_fd, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof opt);
#endif
    }

private:
    short _port;
    int _fd;
};
class tcp_socket_t : public socket_t
{
public:
    virtual ~tcp_socket_t() { ; }
    tcp_socket_t(const uint16_t port)
        : socket_t(port, TCP_SOCKET)
    {
        ;
    }
    // listen
    bool listen_fd(const int num = 128) const
    {
        return listen(get_fd(), num) == 0;
    }
    // accept
    int accept_fd() const
    {
        sockaddr_in peer_addr;
        socketlen_t peer_len = sizeof peer_addr;
        return accept(get_fd(), (sockaddr *)&peer_addr, &peer_len);
    }
};
class udp_socket_t : public socket_t
{
    virtual ~udp_socket_t() { ; }
    udp_socket_t(const uint16_t port)
        : socket_t(port, UDP_SOCKET)
    {
        ;
    };
    int recive(char *buffer, size_t size, sockaddr *peer, socketlen_t *peer_size)
    {
        return recvfrom(get_fd(), buffer, size, 0, peer, peer_size);
    }
    ssize_t send(const char *buffer, const size_t size, const sockaddr *peer, int peer_size)
    {
        return sendto(get_fd(), buffer, size, 0, peer, peer_size);
    }
};
